import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(4, 64),
    nn.ReLU(),
    nn.Linear(64, 128),
    nn.ReLU(),
    nn.Linear(128, 2)
)

# # 1. 保存模型的权重和其他参数
# torch.save(model.state_dict(), 'model_state_dict.pth')
#
# # 2. 保存整个模型
# torch.save(model, 'model.pth')

# 3. 加载模型
net = torch.load('model.pth', weights_only=False)
x = torch.rand(10, 4)
print(net(x))

# 4. 加载模型状态字典
# 状态字典中只保留了模型的权重和其它参数，未保留模型结构
# 你要去用的时候，得先有模型，再将状态字典赋给模型
model.load_state_dict(torch.load('model_state_dict.pth'))